import codecs
import logging
import tarfile
import os
import time
from os import times

from gxl_ai_utils.utils import utils_file
from pygtrans import Translate

# client = Translate(proxies={'https': 'http://localhost:7890'})
client = Translate()

def do_translate(text):
    text = client.translate(text)
    return text.translatedText


def extract_tar_file(tar_path, extract_path):
    # 解压 tar 文件到指定目录
    with tarfile.open(tar_path, "r") as tar:
        tar.extractall(path=extract_path)
def handle_str(text):
    """
    对字符串的具体处理，此处是翻译
    :param text:
    :return:
    """
    return do_translate(text)

def process_txt_files(dir_path):
    num_txt=0
    text_dict = {}
    txt_file_path_list = utils_file.get_file_path_list_for_wav_dir(dir_path, suffix='txt')
    for txt_path in txt_file_path_list:
        key = utils_file.get_file_pure_name_from_path(txt_path)
        text_now=utils_file.load_first_row_clean(txt_path)
        if text_now is None or text_now == '' or len(text_now)< 1:
            continue
        text_dict[key]=text_now
    key_list = list(text_dict.keys())
    value_list = list(text_dict.values())
    chunk_size = 30
    cn_value_list = []
    for i in range(0, len(key_list), chunk_size):
        chunk = value_list[i:i+chunk_size]
        result_list = client.translate(chunk)
        for result in result_list:
            cn_value_list.append(result.translatedText)
        for key, cn_value in zip(key_list, cn_value_list):
            output_path_now = os.path.join(dir_path, key+'.txt')
            with codecs.open(str(output_path_now), 'w', encoding='utf-8') as f:
                    f.write(cn_value)
        num_txt += chunk_size
        print(f'已处理 {num_txt} 个文本文件, 一共{len(key_list)}个文件')



def create_tar_file(output_tar_path, input_dir):
    # 将目录重新压缩为 tar 文件
    with tarfile.open(output_tar_path, "w") as tar:
        tar.add(input_dir, arcname=os.path.basename(input_dir))


# 主函数
def main():
    utils_file.do_config_root_logger()
    for i in range(634):
        print(f'tar 进度：{i}/634')
        tar_num = str(i).zfill(3)
        tar_path = f"/home/node54_tmpdata/xlgeng/caption_data/shards/Audio_Caption_000000{tar_num}.tar"
        extract_path = f"/home/node54_tmpdata/xlgeng/caption_data/shards/Audio_Caption_000000{tar_num}"
        output_tar_path = f"/home/node54_tmpdata/xlgeng/caption_data/shards_2/Audio_Caption_000000{tar_path}.tar"
        if os.path.exists(output_tar_path):
            logging.info(f'{output_tar_path} 已存在，跳过')
            continue
        utils_file.makedir_for_file(output_tar_path)
        # 解压
        extract_tar_file(tar_path, extract_path)
        # 处理 .txt 文件
        process_txt_files(extract_path)
        # 重新打包成 tar 文件
        create_tar_file(output_tar_path, extract_path)


if __name__ == "__main__":
    """"""

    now = time.time()
    main()
    print("cost:", time.time() - now,"s")
    # test translate func
    # input_text1='A rooster crows'
    # input_text2=''
    # input_text_list = ['A rooster crows','A Chinese man is walking']
    # text = client.translate(input_text1)
    # print(text)
    # text = client.translate(input_text2)
    # print(text)
    # text = client.translate(input_text_list)
    # print(text)
